package online.chazi.bootplus.multiple.datasource.router;

import lombok.extern.slf4j.Slf4j;
import online.chazi.bootplus.multiple.datasource.MultipleDataSource;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 说明：按权重
 * Created by 叉子同学 on 2020-05-13 19:58
 */
@Slf4j
public class MultipleDataSourceLoadBalanceWeightRouter implements MultipleDataSourceBalanceRouter {

    //计数, name => weight
    private static final Map<String, Integer> DATA_SOURCE_WEIGHT = new HashMap<>();
    //权重，计数
    private static final Map<String, AtomicInteger> WEIGHT_COUNTER = new HashMap<>();
    //当前group下标
    private static final Map<String, AtomicInteger> GROUP_INDEX = new HashMap<>();

    @Override
    public void initWeight(String name, int weight) {
        String groupName = MultipleDataSource.getGroupName(name);
        //权重
        DATA_SOURCE_WEIGHT.putIfAbsent(name, weight);
        //原子计数器
        WEIGHT_COUNTER.putIfAbsent(groupName, new AtomicInteger(0));
        //分组下标
        GROUP_INDEX.putIfAbsent(groupName, new AtomicInteger(0));
    }

    @Override
    public String determineLookupKey(String ds, List<String> targetDataSources) {
        //计数器
        AtomicInteger counter = WEIGHT_COUNTER.get(ds);
        //下标
        AtomicInteger index = GROUP_INDEX.get(ds);
        //权重
        String name = targetDataSources.get(index.get());
        Integer weight = DATA_SOURCE_WEIGHT.get(name);
        //
        log.info("使用负载均衡。group={}, name={}, counter={}, index={}, weight={}", ds, name, counter.get(), index.get(), weight);
        ds = name;  //变更数据源。
        counter.incrementAndGet();
        if (counter.get() >= weight) {
            counter.set(0);
            //组的下一个
            if (index.get() >= targetDataSources.size() - 1) {
                index.set(0);
            } else {
                index.incrementAndGet();
            }
        }
        return ds;
    }
}
